//=======================================================================
// Base64Encode/Base64Decode
// compliant with RFC 2045
//=======================================================================
//
#ifdef WIN32
#include <Windows.h>
#endif
#include <assert.h>

#include "Base64Util.h"
#include "BaseUtil.h"
#include "MacroUtil.h"

#ifndef INT_MAX
#define INT_MAX       2147483647    /* maximum (signed) int value */
#endif
#ifndef NULL
#define NULL    0
#endif
int Base64EncodeGetRequiredLength(int nSrcLen, unsigned long dwFlags)
{
    long long nSrcLen4=(long long)(nSrcLen)*4;
    assert(nSrcLen/4 <= INT_MAX);

    int nRet = (int)(nSrcLen4/3);

    if ((dwFlags & ATL_BASE64_FLAG_NOPAD) == 0)
        nRet += nSrcLen % 3;

    int nCRLFs = nRet / 76 + 1;
    int nOnLastLine = nRet % 76;

    if (nOnLastLine)
    {
        if (nOnLastLine % 4)
            nRet += 4-(nOnLastLine % 4);
    }

    nCRLFs *= 2;

    if ((dwFlags & ATL_BASE64_FLAG_NOCRLF) == 0)
        nRet += nCRLFs;

    return nRet;
}

int Base64DecodeGetRequiredLength(int nSrcLen)
{
    return nSrcLen;
}

bool Base64Encode(const unsigned char* pbSrcData,
                  int                  nSrcLen,
                  char*                szDest,
                  int*                 pnDestLen,
                  unsigned long        dwFlags)
{
    static const char s_chBase64EncodingTable[64] = {
        'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q',
        'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g',	'h',
        'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y',
        'z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '+', '/' };

    if (!pbSrcData || !szDest || !pnDestLen)
    {
        return false;
    }

    if(*pnDestLen < Base64EncodeGetRequiredLength(nSrcLen, dwFlags))
    {
        assert(false);
        return false;
    }

    int nWritten = 0;
    int nLen1 = (nSrcLen/3)*4 ;
    int nLen2 =nLen1/76 ;
    int nLen3 = 19 ;

    for (int i=0; i<=nLen2; i++)
    {
        if (i==nLen2)
            nLen3 = (nLen1%76)/4;

        for (int j=0; j<nLen3; j++)
        {
            unsigned int dwCurr = 0;
            for (int n=0; n<3; n++)
            {
                dwCurr |= *pbSrcData++;
                dwCurr <<= 8;
            }
            for (int k=0; k<4; k++)
            {
                unsigned char b = (unsigned char)(dwCurr>>26);
                *szDest++ = s_chBase64EncodingTable[b];
                dwCurr <<= 6;
            }
        }
        nWritten+= nLen3*4;

        if ((dwFlags & ATL_BASE64_FLAG_NOCRLF)==0)
        {
            *szDest++ = '\r';
            *szDest++ = '\n';
            nWritten+= 2;
        }
    }

    if (nWritten && (dwFlags & ATL_BASE64_FLAG_NOCRLF)==0)
    {
        szDest-= 2;
        nWritten -= 2;
    }

    nLen2 = (nSrcLen%3) ? (nSrcLen%3 + 1) : 0;
    if (nLen2)
    {
        unsigned int dwCurr = 0;
        for (int n=0; n<3; n++)
        {
            if (n<(nSrcLen%3))
                dwCurr |= *pbSrcData++;
            dwCurr <<= 8;
        }
        for (int k=0; k<nLen2; k++)
        {
            unsigned char b = (unsigned char)(dwCurr>>26);
            *szDest++ = s_chBase64EncodingTable[b];
            dwCurr <<= 6;
        }
        nWritten+= nLen2;
        if ((dwFlags & ATL_BASE64_FLAG_NOPAD)==0)
        {
            nLen3 = nLen2 ? 4-nLen2 : 0;
            for (int j=0; j<nLen3; j++)
            {
                *szDest++ = '=';
            }
            nWritten+= nLen3;
        }
    }

    *pnDestLen = nWritten;
    return true;
}

int DecodeBase64Char(unsigned int ch)
{
    // returns -1 if the character is invalid
    // or should be skipped
    // otherwise, returns the 6-bit code for the character
    // from the encoding table
    if (ch >= 'A' && ch <= 'Z')
        return ch - 'A' + 0;	// 0 range starts at 'A'
    if (ch >= 'a' && ch <= 'z')
        return ch - 'a' + 26;	// 26 range starts at 'a'
    if (ch >= '0' && ch <= '9')
        return ch - '0' + 52;	// 52 range starts at '0'
    if (ch == '+')
        return 62;
    if (ch == '/')
        return 63;
    return -1;
}

bool IsBase64Char(unsigned char ch)
{
    if (ch >= 'A' && ch <= 'Z')
        return true;
    if (ch >= 'a' && ch <= 'z')
        return true;
    if (ch >= '0' && ch <= '9')
        return true;
    if (ch == '+')
        return true;
    if (ch == '/')
        return true;
    return false;
}

bool Base64Decode(const char*    szSrc,
                  int            nSrcLen,
                  unsigned char* pbDest,
                  int*           pnDestLen)
{
    // walk the source buffer
    // each four character sequence is converted to 3 bytes
    // CRLFs and =, and any characters not in the encoding table
    // are skiped

    if (szSrc == NULL || pnDestLen == NULL)
    {
        assert(false);
        return false;
    }

    const char* szSrcEnd = szSrc + nSrcLen;
    int nWritten = 0;

    bool bOverflow = (pbDest == NULL) ? true : false;

    while (szSrc < szSrcEnd &&(*szSrc) != 0)
    {
        unsigned int dwCurr = 0;
        int i;
        int nBits = 0;
        for (i=0; i<4; i++)
        {
            if (szSrc >= szSrcEnd)
                break;
            int nCh = DecodeBase64Char(*szSrc);
            szSrc++;
            if (nCh == -1)
            {
                // skip this char
                i--;
                continue;
            }
            dwCurr <<= 6;
            dwCurr |= nCh;
            nBits += 6;
        }

        if(!bOverflow && nWritten + (nBits/8) > (*pnDestLen))
            bOverflow = true;

        // dwCurr has the 3 bytes to write to the output buffer
        // left to right
        dwCurr <<= 24-nBits;
        for (i=0; i<nBits/8; i++)
        {
            if(!bOverflow)
            {
                *pbDest = (unsigned char) ((dwCurr & 0x00ff0000) >> 16);
                pbDest++;
            }
            dwCurr <<= 8;
            nWritten++;
        }

    }

    *pnDestLen = nWritten;

    if(bOverflow)
    {
        if(pbDest != NULL)
        {
            assert(false);
        }

        return false;
    }

    return true;
}

HRESULT Base64EncodeEx(BYTE*  pbyData,
                       DWORD  dwSize,
                       CHAR** ppszBase64Str,
                       DWORD* pdwBase64Len,
                       DWORD  dwFlags)
{
    HRESULT  hResult = E_FAIL;
    bool     bSuccess = false;
    CHAR*    szBase64Str = NULL;
    INT      nBase64Len = 0;

    ERROR_ESCAPE((!pbyData || dwSize < 1), ERROR_INVALID_PARAMETER, "Check Parameter");

    nBase64Len = Base64EncodeGetRequiredLength(dwSize, dwFlags);
    ALLOCATE_MEMORY(szBase64Str, char, nBase64Len + 1);

    bSuccess = Base64Encode(pbyData, dwSize, szBase64Str, &nBase64Len, dwFlags);
    ERROR_ESCAPE(!bSuccess, ERROR_ERRORS_ENCOUNTERED, "Base64Encode");

    *ppszBase64Str = szBase64Str;
    if (pdwBase64Len)
        *pdwBase64Len = nBase64Len;
    szBase64Str = NULL;

    hResult = S_OK;

exit:
    FREE_MEMORY(szBase64Str);

    return hResult;
}

HRESULT Base64DecodeEx(CCHAR* szBase64Str,
                       BYTE** ppbyDecodedData,
                       DWORD* pdwDecodedSize)
{
    HRESULT hResult = E_FAIL;
    bool    bSuccess = false;
    DWORD   dwBase64Len = 0;
    BYTE*   pbyDecodedData = NULL;
    INT     nDecodedSize = 0;

    ERROR_ESCAPE((!szBase64Str || !ppbyDecodedData || !pdwDecodedSize), ERROR_INVALID_PARAMETER, "Check Parameter");

    dwBase64Len = (DWORD)strlen(szBase64Str);
    nDecodedSize = Base64DecodeGetRequiredLength(dwBase64Len);
    ALLOCATE_MEMORY(pbyDecodedData, BYTE, nDecodedSize);

    bSuccess = Base64Decode(szBase64Str, dwBase64Len, pbyDecodedData, &nDecodedSize);
    ERROR_ESCAPE(!bSuccess, ERROR_INVALID_DATA, "Base64Decode");

    *ppbyDecodedData = pbyDecodedData;
    *pdwDecodedSize = nDecodedSize;
    pbyDecodedData = NULL;

    hResult = S_OK;

exit:
    FREE_MEMORY(pbyDecodedData);

    return hResult;
}